import numpy as np


def apply_augmentations(x_train, y_train, noise_std=0.0, mask_prob=0.0, value_channel_index=0, rng=None):
    """
    Apply simple, safe augmentations to training arrays.

    Parameters
    - x_train: np.ndarray of shape [samples, seq_len, num_nodes, input_dim]
    - y_train: np.ndarray of shape [samples, horizon, num_nodes, output_dim]
    - noise_std: float, Gaussian noise std applied to the value channel
    - mask_prob: float, probability to randomly zero-out the value channel entries
    - value_channel_index: int, which channel in input_dim corresponds to the primary value
    - rng: optional numpy.random.Generator for reproducibility

    Returns
    - (x_train_aug, y_train_aug)
    """
    if rng is None:
        rng = np.random.default_rng()

    x_aug = x_train.copy()
    y_aug = y_train  # labels unchanged for these augmentations

    # Gaussian noise on value channel
    if noise_std and noise_std > 0:
        noise = rng.normal(0.0, noise_std, size=x_aug[..., value_channel_index].shape)
        x_aug[..., value_channel_index] = x_aug[..., value_channel_index] + noise.astype(x_aug.dtype)

    # Random mask/drop on value channel
    if mask_prob and mask_prob > 0:
        mask = rng.random(size=x_aug[..., value_channel_index].shape) < mask_prob
        x_aug[..., value_channel_index][mask] = 0.0

    return x_aug, y_aug